
import cv2
import torch
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

input_location = 'a.png'
input_mask_location = 'a1.png'
prompt = 'background'
output_image_path = './result.png'

input = {
    'image': input_location,
    'mask': input_mask_location,
    'prompt': prompt
}
image_inpainting = pipeline(
    Tasks.image_inpainting,
    model='damo/cv_stable-diffusion-v2_image-inpainting_base',
    device='gpu',
    torch_dtype=torch.float16,
    enable_attention_slicing=True)
output = image_inpainting(input)[OutputKeys.OUTPUT_IMG]
cv2.imwrite(output_image_path, output)
print('pipeline: the output image path is {}'.format(output_image_path))